#!/usr/bin/env python3
"""
分布式训练配置文件
支持多GPU训练
"""

import argparse
import os
import sys
import torch
import torch.distributed as dist
from setup_training import setup_and_train

def parse_args():
    parser = argparse.ArgumentParser(description='分布式训练配置')
    
    # 数据设置
    parser.add_argument('--data_dir', type=str, default='../split_data/webfg400_train', help='数据集路径')
    parser.add_argument('--num_classes', type=int, default=5000, help='类别数量')
    parser.add_argument('--augmentation_config', type=str, default='updated_augmentation_config.yml', help='数据增强配置文件')
    
    # 模型设置
    parser.add_argument('--model_name', type=str, default='SwinTransformerV2', help='模型名称')
    parser.add_argument('--pre_trained_weights_path', type=str, default=None, help='预训练权重路径')
    parser.add_argument('--load_previous_weight_path', type=str, default=None, help='之前训练的权重路径')
    parser.add_argument('--weight_init_type', type=str, default='kaiming_normal', 
                       choices=['xavier_uniform', 'xavier_normal', 'kaiming_uniform', 'kaiming_normal', 'normal', 'uniform'],
                       help='权重初始化方法')
    
    # 训练设置
    parser.add_argument('--batch_size', type=int, default=64, help='批次大小 (每个GPU)')
    parser.add_argument('--epochs', type=int, default=400, help='训练轮数')
    parser.add_argument('--learning_rate', type=float, default=0.00005, help='学习率')
    parser.add_argument('--weight_decay', type=float, default=0.00001, help='权重衰减')
    parser.add_argument('--optimizer', type=str, default='AdamW', help='优化器')
    
    # 学习率调度设置
    parser.add_argument('--scheduler', type=str, default='CosineAnnealingLR', help='学习率调度器')
    parser.add_argument('--T_max', type=int, default=400, help='Maximum number of iterations for CosineAnnealingLR')
    parser.add_argument('--eta_min', type=float, default=1e-6, help='Minimum learning rate for CosineAnnealingLR')
    parser.add_argument('--factor', type=float, default=0.8, help='学习率衰减因子 (ReduceLROnPlateau)')
    parser.add_argument('--patience', type=int, default=20, help='等待轮数 (ReduceLROnPlateau)')
    parser.add_argument('--warmup_epochs', type=int, default=3, help='预热轮数')
    parser.add_argument('--warmup_type', type=str, default='linear', help='预热类型')
    parser.add_argument('--warmup_start_lr', type=float, default=1e-6, help='预热起始学习率')
    parser.add_argument('--enable_warmup', action='store_true', default=True, help='启用预热')
    
    # 损失函数设置
    parser.add_argument('--loss_function', type=str, default='CrossEntropyLoss', help='损失函数')
    
    # 训练策略设置
    parser.add_argument('--grad_clip_norm', type=float, default=1.0, help='梯度裁剪范数')
    parser.add_argument('--check_nan_loss', action='store_true', default=None, help='检查NaN损失')
    parser.add_argument('--use_mixed_precision', action='store_true', default=None, help='使用混合精度训练')
    
    # 日志和保存设置
    parser.add_argument('--log_dir', type=str, default='./logs', help='日志目录')
    parser.add_argument('--weight_dir', type=str, default='./weights', help='权重保存目录')
    parser.add_argument('--random_seed', type=int, default=42, help='随机种子')
    
    # 分布式训练设置
    parser.add_argument('--local_rank', type=int, default=0, help='本地GPU排名')
    parser.add_argument('--world_size', type=int, default=1, help='总进程数')
    parser.add_argument('--rank', type=int, default=0, help='全局排名')
    
    return parser.parse_args()

def setup_distributed():
    """初始化分布式训练环境"""
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        rank = int(os.environ['RANK'])
        world_size = int(os.environ['WORLD_SIZE'])
        local_rank = int(os.environ['LOCAL_RANK'])
        
        # 设置CUDA设备
        torch.cuda.set_device(local_rank)
        
        # 初始化分布式进程组
        dist.init_process_group(backend='nccl')
        
        print(f"🚀 分布式训练初始化完成")
        print(f"   全局排名: {rank}/{world_size}")
        print(f"   本地排名: {local_rank}")
        print(f"   GPU设备: {torch.cuda.current_device()}")
        
        return True, rank, world_size, local_rank
    else:
        print("🔧 单GPU训练模式")
        return False, 0, 1, 0

def cleanup_distributed():
    """清理分布式训练环境"""
    if dist.is_initialized():
        dist.destroy_process_group()

def main():
    # 解析参数
    args = parse_args()
    
    # 初始化分布式训练
    is_distributed, rank, world_size, local_rank = setup_distributed()
    
    # 设置环境变量
    os.environ['LOCAL_RANK'] = str(local_rank)
    os.environ['RANK'] = str(rank)
    os.environ['WORLD_SIZE'] = str(world_size)
    
    try:
        # 调整批次大小（分布式训练时，总批次大小 = batch_size * world_size）
        if is_distributed:
            effective_batch_size = args.batch_size * world_size
            if rank == 0:
                print(f"📊 有效批次大小: {effective_batch_size} (每GPU: {args.batch_size}, GPU数: {world_size})")
        
        # 打印配置信息（只在主进程打印）
        if rank == 0:
            print("=" * 60)
            print("🎯 分布式训练配置")
            print("=" * 60)
            for arg in vars(args):
                print(f"{arg}: {getattr(args, arg)}")
            print("=" * 60)
        
        # 开始训练
        setup_and_train(args)
        
    except Exception as e:
        print(f"❌ 训练失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        # 清理分布式环境
        cleanup_distributed()
    
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
